import torch
from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
import os
from dataset_download import download_file, extract_zip
from data_loader import get_file_paths_and_labels, BirdDataset
from train import train_model, evaluate_model
from model_TCN import TCN

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# 主程序
def main():
    # 数据集相关路径和设置
    # BIRDCLEF_URL = "https://example.com/birdclef2023.zip"  # 替换为实际链接
    BIRDCLEF_URL = "https://ai-studio-online.bj.bcebos.com/v1/b7940294f22c498ea245ccb971d368213cb265d5048048c79a847e0a24993a46?responseContentDisposition=attachment%3Bfilename%3Dbirdclef-2023.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-11-22T08%3A54%3A54Z%2F21600%2F%2F3447afaa726730642c5ba8b11e304dcfa5bca531eafaa10dbd27eed108b7de4a"  # 替换为实际链接
    DATASET_DIR = "./data"
    ZIP_PATH = os.path.join(DATASET_DIR, "birdclef2023.zip")
    EXTRACTED_PATH = os.path.join(DATASET_DIR, "birdclef2023")

    # 确保数据目录存在
    os.makedirs(DATASET_DIR, exist_ok=True)

    # 数据集下载与解压
    if not os.path.exists(EXTRACTED_PATH):
        print("Dataset not found. Downloading and extracting...")
        download_file(BIRDCLEF_URL, ZIP_PATH)
        extract_zip(ZIP_PATH, EXTRACTED_PATH)
    else:
        print(f"Dataset already exists at {EXTRACTED_PATH}.")

    # 加载数据
    print("Loading dataset...")
    train_files, train_labels = get_file_paths_and_labels(os.path.join(EXTRACTED_PATH, "train_audio"))
    val_files, val_labels = get_file_paths_and_labels(os.path.join(EXTRACTED_PATH, "val"))
    test_files, test_labels = get_file_paths_and_labels(os.path.join(EXTRACTED_PATH, "test"))

    train_dataset = BirdDataset(train_files, train_labels)
    val_dataset = BirdDataset(val_files, val_labels)
    test_dataset = BirdDataset(test_files, test_labels)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    # 构建模型
    print("Initializing model...")
    model = TCN(
        input_size=128,  # 梅尔频谱特征维度
        num_classes=train_labels.shape[1],  # 多标签分类的类别数
        num_channels=[128, 256, 256, 128],  # TCN 中的通道数
        kernel_size=3,
        dropout=0.2
    ).to(device)

    # 定义损失函数和优化器
    criterion = nn.BCEWithLogitsLoss()  # 二分类交叉熵损失（适合多标签任务）
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    print("Training model...")
    train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20)

    # 测试模型
    print("Evaluating model...")
    evaluate_model(model, test_loader)


if __name__ == "__main__":
    main()
